package com.saturday.web.filter.xss;

import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.saturday.web.common.PostServletRequest;
import com.saturday.web.utils.XssUtil;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;

import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

/**
 * xss简易防护请求装饰类
 */
public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    /**
     * xss-json值过滤器
     */
    protected XssValueFilter valueFilter;

    /**
     * 是否为加密参数
     */
    private boolean isEncrypt = false;

    public XssHttpServletRequestWrapper(HttpServletRequest request, XssValueFilter valueFilter) {
        super(request);
        this.valueFilter = valueFilter;
    }

    public XssHttpServletRequestWrapper(HttpServletRequest request, XssValueFilter valueFilter, boolean isEncrypt) {
        super(request);
        this.valueFilter = valueFilter;
        this.isEncrypt = isEncrypt;
    }

    @Override
    public String getHeader(String name) {
        return XssCleanUtils.xssClean(super.getHeader(name));
    }

    @Override
    public String getQueryString() {
        return clean(super.getQueryString());
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        return valueFilter.getExcludeAttr().contains(name) ? XssCleanUtils.cleanScript(value) : clean(value);
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values != null) {
            int length = values.length;
            String[] escapseValues = new String[length];
            boolean excludeFlag = valueFilter.getExcludeAttr().contains(name);
            for (int i = 0; i < length; i++) {
                escapseValues[i] = excludeFlag ? XssCleanUtils.cleanScript(values[i]) : clean(values[i]);
            }
            return escapseValues;
        }
        return null;
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> params = new HashMap<>(super.getParameterMap());
        for (Entry<String, String[]> item : params.entrySet()) {
            params.put(item.getKey(), getParameterValues(item.getKey()));
        }
        return params;
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        String responseString = HttpUtil.getString(super.getInputStream(), Charset.defaultCharset(), true);

        if (StringUtils.isNotEmpty(responseString)) {
            Object parse = JSON.parse(responseString);
            return new PostServletRequest((HttpServletRequest) this.getRequest(), JSONObject.toJSONString(processPostBody(parse))).getInputStream();
        }
        return super.getInputStream();
    }

    private Object processPostBody(Object data) {
        if (data instanceof JSONArray) {
            JSONArray jsonArray = (JSONArray) data;
            Object[] objects = jsonArray.stream().toArray();
            for (Object obj : objects) {
                processPostBody(obj);
            }

            return objects;
        } else if (data instanceof JSONObject) {
            Map<String, Object> bodyJson = (Map<String, Object>) data;
            for (Entry<String, Object> entry : bodyJson.entrySet()) {
                String k = entry.getKey();
                Object value = entry.getValue();
                if (value instanceof String) {
                    String cleanValue = XssUtil.clean(value.toString());
                    bodyJson.put(k, cleanValue);
                } else {
                    processPostBody(value);
                }
            }
            return bodyJson;
        }
        return null;
    }

    /**
     * @param value 密文
     * @return 清洗报文
     * @方法名称 clean
     * @功能描述 <pre>解密及xxs字符清洗</pre>
     */
    private String clean(String value) {
        if (StringUtils.isEmpty(value)) {// 密文空值处理
            return null;
        }
        if (StringUtils.isEmpty(value) || value.equals("undefined") || value.equals("null")) {// 原始值处理
            return null;
        }
        if ((value.startsWith("{") || value.startsWith("[")) && XssValueFilter.isJson(value)) {// json走 jsonvalue过滤
            //先过滤，再还原
            return (String) JSON.parse(null != valueFilter ? JSON.toJSONString(value, valueFilter) : JSON.toJSONString(value));
        }
        return XssCleanUtils.xssClean(value);// xss清洗
    }

}
